import logging
import numpy as np
import pandas as pd
import torch
from typing import List, Dict
from datetime import datetime, timedelta
from transformers import AutoModelForCausalLM, AutoTokenizer

from .stock_service import StockService
from ..models.prediction_models import (
    PredictionRequest, PredictionResponse, PredictionError,
    PredictionPoint, ConfidenceInterval
)

logger = logging.getLogger(__name__)


class PredictionService:
    """使用 Kronos-small 模型的股票价格预测服务"""

    def __init__(self):
        self.stock_service = StockService()
        self.predictor = None
        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
        self._model_loaded = False

    async def load_model(self):
        """加载 Kronos-small 模型"""
        if self._model_loaded:
            return

        try:
            logger.info("正在加载 Kronos 模型: NeoQuasar/Kronos-small")

            from model import Kronos, KronosTokenizer, KronosPredictor

            # KronosTokenizer 初始化参数
            tokenizer = KronosTokenizer(
                d_in=6,  # 股票数据维度 (open, high, low, close, volume, amount)
                d_model=512,  # 模型维度
                n_heads=8,  # 注意力头数
                ff_dim=2048,  # 前馈网络维度
                n_enc_layers=6,  # 编码器层数
                n_dec_layers=6,  # 解码器层数
                ffn_dropout_p=0.1,  # 前馈网络 dropout
                attn_dropout_p=0.1,  # 注意力层 dropout
                resid_dropout_p=0.1,  # 残差连接 dropout
                s1_bits=8,  # S1 量化位数
                s2_bits=8,  # S2 量化位数
                beta=0.25,  # BSQuantizer beta 参数
                gamma0=0.9,  # BSQuantizer gamma0 参数
                gamma=0.99,  # BSQuantizer gamma 参数
                zeta=1e-4,  # BSQuantizer zeta 参数
                group_size=4  # BSQuantizer 组大小
            )
            
            # Kronos 模型初始化参数
            model = Kronos(
                s1_bits=8,  # S1 量化位数
                s2_bits=8,  # S2 量化位数
                n_layers=12,  # Transformer 层数
                d_model=512,  # 模型维度
                n_heads=8,  # 注意力头数
                ff_dim=2048,  # 前馈网络维度
                ffn_dropout_p=0.1,  # 前馈网络 dropout
                attn_dropout_p=0.1,  # 注意力层 dropout
                resid_dropout_p=0.1,  # 残差连接 dropout
                token_dropout_p=0.1,  # token dropout
                learn_te=True  # 学习时间嵌入
            )
            
            self.predictor = KronosPredictor(model, tokenizer, device=self.device, max_context=512)

            self._model_loaded = True
            logger.info("Kronos 模型加载成功")

        except Exception as e:
            logger.error(f"Kronos 模型加载失败: {str(e)}")
            raise PredictionError(f"模型加载失败: {str(e)}")

    async def predict_stock_price(self, request: PredictionRequest) -> PredictionResponse:
        """执行股票价格预测"""
        try:
            logger.info(f"开始预测股票 {request.stock_code}，预测天数: {request.prediction_days}，历史数据天数: {request.historical_days}")

            await self.load_model()
            historical_data = await self._get_historical_data(request.stock_code, request.historical_days)

            # 转换为 DataFrame
            df = pd.DataFrame(historical_data)
            df['timestamps'] = pd.to_datetime(df['date'])
            x_df = df[['open', 'high', 'low', 'close']].copy()
            x_ts = df['timestamps']

            # 构造预测时间戳
            last_date = x_ts.iloc[-1]
            y_dates = []
            while len(y_dates) < request.prediction_days:
                last_date += timedelta(days=1)
                if last_date.weekday() < 5:  # 跳过周末
                    y_dates.append(last_date)
            y_ts = pd.Series(y_dates)

            # 使用 Kronos 预测
            pred_df = self.predictor.predict(
                df=x_df,
                x_timestamp=x_ts,
                y_timestamp=y_ts,
                pred_len=request.prediction_days,
                T=1.0,
                top_p=0.9,
                sample_count=1
            )

            # 格式化预测结果
            prediction_points = []
            for date, row in pred_df.iterrows():
                prediction_points.append(
                    PredictionPoint(
                        date=date.strftime('%Y-%m-%d'),
                        predicted_close=round(float(row['close']), 2)
                    )
                )

            # 置信区间（仍用启发式，5% 起步，每天+2%）
            confidence_intervals = []
            for i, point in enumerate(prediction_points):
                factor = 0.05 + i * 0.02
                margin = point.predicted_close * factor
                confidence_intervals.append(
                    ConfidenceInterval(
                        date=point.date,
                        lower_bound=max(0, point.predicted_close - margin),
                        upper_bound=point.predicted_close + margin
                    )
                )

            return PredictionResponse(
                stock_code=request.stock_code,
                prediction_days=request.prediction_days,
                predictions=prediction_points,
                confidence_intervals=confidence_intervals,
                historical_data_count=len(historical_data),
                historical_data=historical_data
            )

        except Exception as e:
            logger.error(f"预测服务错误: {str(e)}")
            raise PredictionError(f"预测服务错误: {str(e)}")

    async def _get_historical_data(self, stock_code: str, historical_days: int = 100) -> List[Dict]:
        """获取历史股票数据"""
        end_date = datetime.now()
        # 为了确保获取到足够的交易日数据，我们获取更多天数（考虑周末和节假日）
        calendar_days = int(historical_days * 1.5)  # 增加50%的缓冲
        start_date = end_date - timedelta(days=calendar_days)
        
        result = await self.stock_service.get_historical_data(
            stock_code,
            start_date.strftime('%Y-%m-%d'),
            end_date.strftime('%Y-%m-%d')
        )
        if not result or not result.get('data'):
            raise PredictionError(f"股票 {stock_code} 没有可用的历史数据")
        
        # 取最近的指定天数的数据
        data = result['data']
        if len(data) > historical_days:
            data = data[-historical_days:]
        
        logger.info(f"获取到 {len(data)} 天的历史数据用于预测")
        return data
